Greptile Summary

This PR fixes a bug in THD-format MLA (e.g. DeepSeek V3) where FlashAttention 2 was entirely blocked for mismatched Q/K and V head dimensions. It introduces zero-padding of V (and optionally Q/K) up to max(head_dim_qk, head_dim_v) before the FA2 call and trims the output back to the original V head dimension afterward, while also tightening the FA2 head-dim validity check to account for the padded dimension.

• Adds _pad_qkv_head_dim and _trim_output helpers; applies them in the FA2 branch when head_dim_qk != head_dim_v and the backend is FA2 (guards correctly skip FA3/FA4 which support MLA natively).
• Removes the blanket FA2-disable guard for mismatched head dims in utils.py and replaces it with a fa2_padded_head_dim-based validity check that also enforces a >192 restriction on older architectures.

Confidence Score: 3/5

The non-FP8 THD MLA path works correctly, but the removed FA2 guard combined with the Float8TensorStorage exclusion from padding leaves the Float8 + MLA + FA2 combination unprotected — FA2 would receive tensors with mismatched head dimensions.

The removed blanket FA2 guard for head_dim_qk != head_dim_v in utils.py is not fully compensated by the padding logic in dot_product_attention.py, which skips Float8TensorStorage inputs. If a Float8 MLA configuration reaches FA2, it will call FA2 with unpadded mismatched head dims, causing a crash or incorrect results. The same guard previously covered this case safely.

Both changed files interact to create the regression: utils.py removes the guard that blocked FA2 for all mismatched-head-dim cases, while dot_product_attention.py adds padding but excludes Float8 tensors.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["DotProductAttention.forward - MLA head_dim_qk != head_dim_v"] --> B{use_flash_attention?}
    B -- No --> C[FusedAttention or Unfused]
    B -- Yes --> D{backend == FA2 version?}
    D -- No --> E[FA3/FA4 support MLA natively - no padding needed]
    D -- Yes --> F{value is Float8TensorStorage?}
    F -- Yes --> G["Skip padding - FA2 receives mismatched head dims - potential crash"]
    F -- No --> H[_pad_qkv_head_dim - pad V to head_dim_qk]
    H --> I[flash_attention with padded Q/K/V]
    I --> J{orig_qk_dim > orig_v_dim?}
    J -- Yes --> K[_trim_output - slice back to orig_head_dim_v]
    J -- No --> L[Return attn_out as-is]
    K --> M[Correct output]
    L --> M

_{Reviews (9): Last reviewed commit: "Support when v is larger than qk" | Re-trigger Greptile}

This change should only be required by the FlashAttention backend. The other two backends FusedAttention and UnfusedDPA do support MLA (head_dim_qk != head_dim_v). I'd propose a few changes:

• move the F.pad() code to if use_flash_attention: branch, and only call _trim_thd_output in that branch
• move _trim_thd_output to the top of the dpa file
• enable the MLA support for FlashAttention:

  TransformerEngine/transformer_engine/pytorch/attention/dot_product_attention/utils.py

  Line 706 in 3c62f42

  if head_dim_qk != head_dim_v:

• remove the newly added test since it'll be covered by the test_dpa_qkv_layout_thd test already once the change in utils.py is done:

  TransformerEngine/tests/pytorch/attention/test_attention.py

  Line 1021 in 3c62f42

  def test_dpa_qkv_layout_thd(dtype, model_configs, model, qkv_layout):

@vcherepanov-nv, could you help push this PR through the finish line? Thanks!

Thank you @cyanguwa, I just cleaned up the PR and also follow your requirements. Please let me know what you think @vcherepanov-nv.

/te-ci pytorch L0

/te-ci pytorch L0

Thank you for pointing this out @cyanguwa, originally I didn't handle this v > qk as this is not a practice for MLA, but since test cases cover this, I have just pushed the changes accordingly.

/te-ci pytorch L0